import sys
import argparse
import numpy as np
import paddle
from paddle import fluid
import matplotlib
import matplotlib.pyplot as plt


def parse_args():
    """
    参数初始化
    :return:
    """
    parse = argparse.ArgumentParser('./data/paddle/fit_a_line')
    parse.add_argument('--enable_ce', action='store_true', help='If set, run the task with continuous evaluation logs.')
    parse.add_argument('--enable_gpu', type=bool, default=False, help='Whether to use GPU or not')
    parse.add_argument('--num_epochs', type=int, default=100, help='number of epochs.')
    return parse.parse_args()


def train_test(executor, program, reader, feeder, fetch_list):
    """
    for training test cost

    :param executor:
    :param program:
    :param reader:
    :param feeder:
    :param fetch_list:
    :return:
    """
    accumulated = 1 * [0]
    count = 0
    for data_test in reader():
        outs = executor.run(program=program, feed=feeder.feed(data_test), fetch_list=fetch_list)
        accumulated = [x[0] + x[1][0] for x in zip(accumulated, outs)]
        count += 1
    # 准确率
    return [x_d / count for x_d in accumulated]


def plot_result(points1, points2):
    """
    真实值与预测值画图呈现
    :param points1:
    :param points2:
    :return:
    """
    matplotlib.use('Agg')
    x1 = [idx for idx in range(len(points1))]
    y1 = points1
    y2 = points2
    l1 = plt.plot(x1, y1, 'r--', label='predict')
    l2 = plt.plot(x1, y2, 'g--', label='true')
    plt.plot(x1, y1, 'ro--', x1, y2, 'g+-')
    plt.title('predict vs true')
    plt.legend()
    plt.show()


def load_data(filename, feature_num=14, ratio=0.8):
    """
    读取文件数据，标准化以及切分训练集和测试集

    :param filename:
    :param feature_num:
    :param ratio:
    :return:
    """
    # 读出文件全部数据，按sep split后为list
    data = np.fromfile(filename, sep=' ')
    # reshape成506行14列，恢复原始的数据格式
    data = data.reshape(data.shape[0] // feature_num, feature_num)
    # 每列（axis=0）极值及均值
    maximums, minimums, avgs = data.max(axis=0), data.min(axis=0), data.sum(axis=0) / data.shape[0]

    # feature_range(maximums[:-1], minimums[:-1])
    # 列值进行标准化
    for i in range(feature_num - 1):
        data[:, i] = (data[:, i] - avgs[i]) / (maximums[i] - minimums[i])
    # 按比例切分训练集和测试集
    offset = int(data.shape[0] * ratio)

    def read_train():
        # 训练集迭代器
        for d in data[: offset]:
            yield d[:-1], d[-1:]

    def read_test():
        # 测试集迭代器
        for d in data[offset:]:
            yield d[:-1], d[-1:]

    return read_train, read_test


def run():
    uci_housing_data = './data/paddle/uci_housing/housing.data'
    batch_size = 20

    train_data, test_data = load_data(uci_housing_data)

    args = parse_args()
    if args.enable_ce:
        train_data_set = train_data
        test_data_set = test_data
    else:
        train_data_set = paddle.reader.shuffle(train_data, buf_size=500)
        test_data_set = paddle.reader.shuffle(test_data, buf_size=500)

    # reader batch
    train_reader = paddle.batch(train_data_set, batch_size=batch_size)
    test_reader = paddle.batch(test_data_set, batch_size=batch_size)

    # 类似tf？
    x = fluid.layers.data(name='x', shape=[13], dtype='float32')
    y = fluid.layers.data(name='y', shape=[1], dtype='float32')

    main_program = fluid.default_main_program()
    startup_program = fluid.default_startup_program()

    if args.enable_ce:
        main_program.random_seed = 90
        startup_program.random_seed = 90

    y_predict = fluid.layers.fc(input=x, size=1, act=None)
    cost = fluid.layers.square_error_cost(input=y_predict, label=y)
    avg_loss = fluid.layers.mean(cost)

    test_program = main_program.clone(for_test=True)

    # 优化器
    sgd_optimizer = fluid.optimizer.SGD(learning_rate=0.001)
    sgd_optimizer.minimize(avg_loss)

    # executor初始化
    place = fluid.CPUPlace()
    exe = fluid.Executor(place)

    params_dirname = './data/paddle/fit_a_line_model'
    num_epochs = args.num_epochs

    # train loop
    feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
    exe.run(startup_program)

    train_prompt = 'Train cost'
    test_prompt = 'Test cost'
    step = 0

    exe_test = fluid.Executor(place)

    # 模型训练
    for pass_id in range(num_epochs):
        # 一次训练
        for data_train in train_reader():
            avg_loss_value, = exe.run(main_program, feed=feeder.feed(data_train), fetch_list=[avg_loss])
            if step % 10 == 0:
                print('{}, step: {}, cost: {}'.format(train_prompt, step, avg_loss_value[0]))
            # 每100步测试一次模型准确率，达到准确率/循环结束 即退出训练
            if step % 100 == 0:
                test_metics = train_test(executor=exe_test, program=test_program, reader=test_reader,
                                         fetch_list=[avg_loss], feeder=feeder)
                print('{}, step: {}, cost: {}'.format(test_prompt, step, test_metics[0]))
                # accuracy 达到阈值，退出训练
                if test_metics[0] < 10.0:
                    break
            step += 1

            if np.isnan(float(avg_loss_value[0])):
                sys.exit("got nan loss, training failed.")

        # 保存模型
        if params_dirname is not None:
            fluid.io.save_inference_model(params_dirname, ['x'], [y_predict], exe)

        if args.enable_ce and pass_id == args.num_epochs - 1:
            print('kpis train_cost {}'.format(avg_loss_value[0]))
            print('kpis test_cost {}'.format(test_metics[0]))

    # 模型预测
    infer_exe = fluid.Executor(place)
    inference_scope = fluid.core.Scope()

    with fluid.scope_guard(inference_scope):
        # load模型
        [inference_program, feed_target_names, fetch_targets] = fluid.io.load_inference_model(params_dirname, infer_exe)
        batch_size = 10

        infer_reader = paddle.batch(test_data, batch_size=batch_size)
        infer_data = next(infer_reader())
        infer_feat = np.array([data[0] for data in infer_data]).astype('float32')
        infer_label = np.array([data[1] for data in infer_data]).astype('float32')

        assert feed_target_names[0] == 'x'
        # 预测结果
        results = infer_exe.run(inference_program,
                                feed={feed_target_names[0]: np.array(infer_feat)},
                                fetch_list=fetch_targets)

        for idx, val in enumerate(results[0]):
            true_value = infer_label[idx][0]
            predict_value = val[0]
            print('{} 真实值： {}， 预测值： {}, 偏差： {}'.format(idx, true_value, predict_value,
                                                       (true_value - predict_value) / true_value))

        plot_result(results[0], infer_label)


if __name__ == '__main__':
    run()
